package com.chatplus.application.aiprocessor.listener.impl;

import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.json.JSONUtil;
import com.chatplus.application.aiprocessor.enumeration.MessageTypeEnum;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.aiprocessor.listener.StreamEventSourceListener;
import com.chatplus.application.enumeration.AiPlatformEnum;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import com.unfbx.chatgpt.entity.chat.FunctionCall;
import com.unfbx.chatgpt.entity.chat.Message;
import com.unfbx.chatgpt.entity.completions.Completion;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;

/**
 * OpenAI监听器
 */
public class ChatGPTEventSourceListener extends StreamEventSourceListener {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(ChatGPTEventSourceListener.class);
    private String reply = "";
    // 函数参数
    private String arguments = "";
    // 函数名称
    private String name = "";
    public String getArguments() {
        return arguments;
    }

    public String getName() {
        return name;
    }

    private final CountDownLatch countDownLatch;

    public ChatGPTEventSourceListener(WebSocketSession session, AiPlatformEnum platform, String prompt,CountDownLatch countDownLatch) {
        super(session, platform, prompt);
        this.countDownLatch = countDownLatch;
    }
    @Override
    public void onEvent(@NotNull EventSource eventSource, String id, String type, @NotNull String data) {
        LOGGER.message("OpenAI返回数据").context("data", data).info();
        try {
            if (data.equals("[DONE]")) {
                LOGGER.message("回答完成").info();
                if (StringUtils.isEmpty(name)) {
                    Completion completion = Completion.builder().prompt(reply).build();
                    chatAiUtilHandler.saveChatHistory(reply, completion.tokens());
                    chatAiUtilHandler.sendMessage(MessageTypeEnum.WS_END.getValue());
                }
                countDownLatch.countDown();
                return;
            }
            ObjectMapper objectMapper = SpringUtil.getBean(ObjectMapper.class);

            ChatCompletionResponse chatCompletionResponse = JSONUtil.toBean(data, ChatCompletionResponse.class);
            FunctionCall functionCall = chatCompletionResponse.getChoices().getFirst().getDelta().getFunctionCall();
            if (Objects.nonNull(functionCall)) {
                if (StringUtils.isNotEmpty(functionCall.getName())) {
                    name += functionCall.getName();
                }
                arguments += functionCall.getArguments();
            }
            ChatCompletionResponse completionResponse = objectMapper.readValue(data, ChatCompletionResponse.class); // 读取Json
            Message delta = completionResponse.getChoices().getFirst().getDelta();
            if (Objects.nonNull(delta.getContent())) {
                chatAiUtilHandler.sendMessage(delta.getContent());
                reply += delta.getContent();
            }
        } catch (IOException e) {
            LOGGER.message("输出回答异常").exception(e).error();
            chatAiUtilHandler.sendMessage(MessageTypeEnum.WS_END.getValue());
            countDownLatch.countDown();
        }

    }
    @Override
    public void onClosed(EventSource eventSource) {
        super.onClosed(eventSource);
        countDownLatch.countDown();
    }

    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        super.onFailure(eventSource, t, response);
        countDownLatch.countDown();
    }
}
